#!/bin/bash
# Copyright 2024. Huawei Technologies Co.,Ltd. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================

cur_path=$(dirname "$(readlink -f "$0")")

so_path=$1
mx_rec_package_path=$2
hccl_cfg_json=$3
dlrm_criteo_data_path=$4

export RANK_SIZE=8
echo "RANK_SIZE=${RANK_SIZE}, please make sure hccl configuration json file match this parameter"
export RANK_TABLE_FILE=${hccl_cfg_json}

################# 参数配置 ######################
export USE_DYNAMIC=0            # 0：静态shape；1：动态shape
export CACHE_MODE="HBM"         # HBM；DDR；SSD
export USE_FAAE=0               # 0：关闭准入淘汰；1：开启准入淘汰
export USE_DYNAMIC_EXPANSION=0  # 0：关闭动态扩容；1: 开启动态扩容
export USE_MULTI_LOOKUP=0       # 0：一表一查；1：一表多查
export USE_MODIFY_GRAPH=0       # 0：feature spec模式；1：自动改图模式
################################################

echo "CACHE_MODE:${CACHE_MODE}"
if [ ${CACHE_MODE} = "SSD" ]; then
  echo "SSD train mode not allow file exist before training,
        deleting dir ${cur_path}/ssd_data then create for SSD use case"
  rm -rf ssd_data
  mkdir ssd_data
fi

export HCCL_CONNECT_TIMEOUT=1200

export DLRM_CRITEO_DATA_PATH=${dlrm_criteo_data_path}
export PYTHONPATH=${mx_rec_package_path}:${so_path}:$PYTHONPATH
export LD_PRELOAD=/usr/lib64/libgomp.so.1
export LD_LIBRARY_PATH=${so_path}:/usr/local/lib:$LD_LIBRARY_PATH

rm -rf kernel*
rm -rf /root/ascend/log/*
rm -rf model_dir_rank* op_cache

export ASCEND_DEVICE_ID=0
export RANK_ID_START=0
export JOB_ID=10086
export CUSTOMIZED_OPS_LIB_PATH=${so_path}/libcust_ops.so # Todo: please config
export MXREC_LOG_LEVEL="INFO"
export TF_CPP_MIN_LOG_LEVEL=3
export ASCEND_GLOBAL_LOG_LEVEL=3
#export USE_FAAE=1
export ENABLE_FORCE_V2_CONTROL=1

export PROFILING_OPTIONS='{"output":"/home/yz/profiling",
                           "training_trace":"on",
                           "task_trace":"on",
                           "aicpu":"on",
                           "fp_point":"",
                           "bp_point":"",
                           "aic_metrics":"PipeUtilization"}'

RANK_ID_START=0

export MXREC_MODE="ASC"
echo "MXREC_MODE is $MXREC_MODE"
export USE_MPI=1
echo "USE_MPI is $USE_MPI"
export py=main_mxrec.py
echo "py is $py"


if [ $USE_MPI -eq 0 ]; then
  echo "use for loop to start tasks"
  for((RANK_ID=$RANK_ID_START;RANK_ID<$((RANK_SIZE+RANK_ID_START));RANK_ID++));
  do
      #设置环境变量，不需要修改
      echo "Device ID: $RANK_ID"
      export RANK_ID=$RANK_ID
      export ASCEND_DEVICE_ID=$RANK_ID
      ASCEND_DEVICE_ID=$RANK_ID
    if [   -d $cur_path/output/${ASCEND_DEVICE_ID} ];then
       rm -rf $cur_path/output/${ASCEND_DEVICE_ID}
       mkdir -p $cur_path/output/${ASCEND_DEVICE_ID}
    else
       mkdir -p $cur_path/output/${ASCEND_DEVICE_ID}
    fi
      nohup python3 ${py} > $cur_path/output/$ASCEND_DEVICE_ID/test_$ASCEND_DEVICE_ID.log 2>&1 &
  done
else
  echo "use horovod to start tasks"
  # GLOG_stderrthreshold -2:TRACE -1:DEBUG 0:INFO 1:WARN 2.ERROR, 默认为INFO
  mpi_args='-x BIND_INFO="0:12 12:48 60:48" -x GLOG_stderrthreshold=2 -x GLOG_logtostderr=true -bind-to none -x NCCL_SOCKET_IFNAME=docker0 -mca btl_tcp_if_exclude docker0'
  interface="lo"

  horovodrun --network-interface ${interface} -np ${RANK_SIZE} --mpi-args "${mpi_args}" --mpi -H localhost:${RANK_SIZE} \
    python3.7 ${py} 2>&1 | tee temp_${CACHE_MODE}_${RANK_SIZE}p.log
fi


